# encoding: utf-8
__author__ = 'Gui'
"""
@Time:2022年6月5日 18:15
@Auth:19级机器人工程 (按拼音排序) 桂源泽 苏琦 颜欢
@File:fira_run.py
@IDE:fira项目自动驾驶主程序
@Software: PyCharm
"""

import rospy
import cv2
import os
import time
import math
import numpy as np
import tensorflow as tf
from robot import Robot
from geometry_msgs.msg import Twist

# 1:[1,0,0,0,0] 前
# 2:[0,1,0,0,0] 左
# 3:[0,0,1,0,0] 右
# 4:[0,0,0,1,0] 后
# 5:[0,0,0,0,1] 冲

# 自动驾驶参数：
threshold_r = 50  # 红二值化阈值 110
threshold_b = 40  # 蓝二值化阈值 110
threshold_g = 30  # 绿二值化阈值 110
flag_max = 5  # 数组判断最大值
rate = 10
# v = 0.6 # 前进速度
v0 = 0.5
angular = 0.4  # 旋转角度
# angular0 = 0.4
vindex = 0.3
aindex = -0.2


# 图像处理参数
width = 480
height = 180
channel = 1
temp_image = np.zeros(width * height * channel, 'uint8')
cap = cv2.VideoCapture(0)


# tf参数
inference_path = tf.Graph()
filepath = os.getcwd() + '/model/3box300/-82'
# /number is model name


def auto_pilot():  # 自主前进程序

    # 机器人初始化
    robot = Robot()
    rate_run = rospy.Rate(rate)
    twist = Twist()

    # tensorflow初始化
    with tf.Session(graph=inference_path) as sess:
        init = tf.global_variables_initializer()
        sess.run(init)
        saver = tf.train.import_meta_graph(filepath + '.meta')  # 调用训练的模型
        saver.restore(sess, filepath)

        tf_X = sess.graph.get_tensor_by_name('input:0')  # 调用所需要的参数
        pred = sess.graph.get_operation_by_name('pred')
        number = pred.outputs[0]
        prediction = tf.argmax(number, 1)

        time_start = time.time()  # 定义开始时间

        while True:  # 开始
            # 定义速度衰减
            time_now = time.time()
            time_v = time_now-time_start
            v = v0*math.exp(vindex*time_v)
            # angular = angular0*math.exp(aindex*time_v)

            frame = robot.get_image()
            resized_height = int(width * 0.75)  # 定义画面新高度为480*0.75

            frame = cv2.resize(frame, (width, resized_height))  # 画面分辨率调整
            # frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)  # 画面转化为灰度

            # _, frame = cv2.threshold(frame, threshold_yuzhi, 255, cv2.THRESH_BINARY) # 二值化
            res = frame[90:270, :]  # 选取摄像头下半部分
            [aisle_b, aisle_g, aisle_r] = cv2.split(res)

            _, aisle_g = cv2.threshold(aisle_g, threshold_g, 255, cv2.THRESH_BINARY) # 统二值化
            aisle_b = cv2.adaptiveThreshold(aisle_b,255,cv2.ADAPTIVE_THRESH_MEAN_C,cv2.THRESH_BINARY,3,5)
            _, aisle_r = cv2.threshold(aisle_r, threshold_r, 255, cv2.THRESH_BINARY) # 统二值化

            res_tmp = cv2.vconcat([aisle_b,aisle_g,aisle_r])
            res = cv2.resize(res_tmp, (width, height), interpolation=cv2.INTER_AREA)
            cv2.imshow("frame", res)  # 显示res图像画面
            cv2.waitKey(1)

            frame = np.array(res, dtype=np.float32)  # 将res数据转化成numpy数组形式
            # 进入机器人运动判断：
            flag = [0, 0, 0, 0, 0]  # 判断次数数组，每一位代表判断的类型

            while max(flag) < flag_max:  # 每一位最大
                # 判断模型：
                value = prediction.eval(feed_dict={tf_X: np.reshape(frame, [-1, height, width, channel])})  # 预测当前画面模型
                if value == 0:
                    flag[0] += 1
                elif value == 1:
                    flag[1] += 1
                elif value == 2:
                    flag[2] += 1
                elif value == 3:
                    flag[3] += 1
                elif value == 4:
                    flag[4] += 1
                elif cv2.waitKey(1) & 0xFF == ord('q'):
                    break

            value_true = np.argmax(np.array(flag))  # 返回最大数值索引值，即预测最大的位置
            if value_true == 0:
                print("forward")
                twist.linear.x = v
                twist.angular.z = 0
            elif value_true == 1:
                print("left")
                twist.linear.x = v
                twist.angular.z = angular
            elif value_true == 2:
                print("right")
                twist.linear.x = v
                twist.angular.z = -angular
            elif value_true == 3:
                print("right")
                twist.linear.x = 0
                twist.angular.z = 0
            elif value_true == 4:
                print("rush!!!!")
                twist.linear.x = 1.2*v
                twist.angular.z = 0

            elif cv2.waitKey(1) & 0xFF == ord('q'):
                break

            # 发布话题消息
            robot.publish_twist(twist)
            rate_run.sleep()

            # flag = [0, 0, 0, 0, 0]

        cv2.destroyAllWindows()


if __name__ == '__main__':
    x = input("################waiting for begin####################")
    auto_pilot()


